from zipline.api import symbol, order_target_percent, record, get_datetime
from zipline import run_algorithm
import pandas as pd
import numpy as np

def initialize(context):
    # 策略参数
    context.rsi_period = 14
    context.overbought = 70
    context.oversold = 30
    context.asset = symbol('AAPL')  # 示例资产
    
def handle_data(context, data):
    # 获取历史价格数据
    prices = data.history(context.asset, 'price', context.rsi_period + 1, '1d')
    
    # 计算价格变化
    deltas = np.diff(prices)
    
    # 计算收益和损失
    gains = np.where(deltas > 0, deltas, 0)
    losses = np.where(deltas < 0, -deltas, 0)
    
    # 计算平均收益和平均损失
    avg_gain = np.mean(gains[-context.rsi_period:])
    avg_loss = np.mean(losses[-context.rsi_period:])
    
    # 计算RSI
    if avg_loss == 0:
        rsi = 100
    else:
        rs = avg_gain / avg_loss
        rsi = 100 - (100 / (1 + rs))
    
    # 记录RSI值
    record(rsi=rsi)
    
    # 交易逻辑
    if rsi < context.oversold:
        order_target_percent(context.asset, 1.0)  # 全仓买入
    elif rsi > context.overbought:
        order_target_percent(context.asset, -1.0)  # 全仓卖空
    else:
        order_target_percent(context.asset, 0.0)  # 平仓

# 示例回测配置
if __name__ == '__main__':
    start = pd.Timestamp('2020-01-01', tz='utc')
    end = pd.Timestamp('2023-01-01', tz='utc')
    
    results = run_algorithm(
        start=start,
        end=end,
        initialize=initialize,
        handle_data=handle_data,
        capital_base=10000,
        data_frequency='daily',
        bundle='quandl'
    )